/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.druid.query.scan;

import org.apache.druid.java.util.common.UOE;
import org.apache.druid.segment.column.ColumnHolder;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ScanQueryTestHelper
{
  public static ScanResultValue generateScanResultValue(
      long timestamp,
      ScanQuery.ResultFormat resultFormat,
      int batchSize
  )
  {
    String segmentId = "some_segment_id";
    List<String> columns = new ArrayList<>(Arrays.asList(ColumnHolder.TIME_COLUMN_NAME, "name", "count"));
    List<Object> events = new ArrayList<>();
    for (int i = 0; i < batchSize; i++) {
      Object event;
      if (resultFormat.equals(ScanQuery.ResultFormat.RESULT_FORMAT_LIST)) {
        Map<String, Object> eventMap = new HashMap<>();
        eventMap.put(ColumnHolder.TIME_COLUMN_NAME, timestamp);
        eventMap.put("name", "Feridun");
        eventMap.put("count", i);
        event = eventMap;
      } else if (resultFormat.equals(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)) {
        event = new ArrayList<>(Arrays.asList(
            timestamp,
            "Feridun",
            i
        ));
      } else {
        throw new UOE("Result format [%s] not supported yet", resultFormat.toString());
      }
      events.add(event);
    }
    return new ScanResultValue(segmentId, columns, events);
  }

  public static List<Map<String, Object>> getEventsListResultFormat(ScanResultValue scanResultValue)
  {
    return (List<Map<String, Object>>) scanResultValue.getEvents();
  }

  public static List<List<Object>> getEventsCompactedListResultFormat(ScanResultValue scanResultValue)
  {
    return (List<List<Object>>) scanResultValue.getEvents();
  }

}
